# -*- coding: utf-8 -*-
"""
相比于 YOLO V3，主干特征提取网络由 DarkNet53 -> CSPDarkNet53
CSP 是可以增加 CNN 学习能力的新型网络结构，CSPNet 将底层的特征映射分为两部分：
1. 一部分经过密集块和过渡层
2. 另一部分与传输的特征映射结合到下一阶段
@author: libo
"""
import torch
import torch.nn.functional as F
import torch.nn as nn
import math
from collections import OrderedDict

class Mish(nn.Module):
    """ MISH激活函数 """
    def __init__(self):
        super(Mish, self).__init__()

    def forward(self, x):
        return x * torch.tanh(F.softplus(x))


class BasicConv(nn.Module):
    """ 卷积块 CONV + BATCHNORM + MISH """
    def __init__(self, in_channels, out_channels, kernel_size, stride=1):
        super(BasicConv, self).__init__()

        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, kernel_size//2, bias=False)
        self.bn = nn.BatchNorm2d(out_channels)
        self.activation = Mish()

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.activation(x)
        return x


class Resblock(nn.Module):
    """ CSPdarknet的结构块的组成部分  内部堆叠的残差块 """
    def __init__(self, channels, hidden_channels=None, residual_activation=nn.Identity()):
        super(Resblock, self).__init__()

        if hidden_channels is None:
            hidden_channels = channels

        self.block = nn.Sequential(
            BasicConv(channels, hidden_channels, 1),
            BasicConv(hidden_channels, channels, 3)
        )

    def forward(self, x):
        return x + self.block(x)


class Resblock_body(nn.Module):
    """ CSPDarknet的结构块，存在一个大残差边，这个大残差边绕过了很多的残差结构
    这部分由一次下采样和多次残差结构的堆叠构成，Darknet53 便是由 resblock_body 模块组合而成；
    相比于 YOLO_v3，这部分进行了一定的修改：
    1. 将 DarknetConv2D 的激活函数由 LeakyReLU 修改成了 Mish: DarknetConv2D_BN_Leaky -> DarknetConv2D_BN_Mish
    2. 修改了 resblock_body 的结构，使用了 CSPNet 结构
    """
    def __init__(self, in_channels, out_channels, num_blocks, first):
        super(Resblock_body, self).__init__()

        """ BasicConv: 卷积块 CONV + BATCHNORM + MISH """
        self.downsample_conv = BasicConv(in_channels, out_channels, 3, stride=2)

        if first:
            self.split_conv0 = BasicConv(out_channels, out_channels, 1)
            self.split_conv1 = BasicConv(out_channels, out_channels, 1)  
            self.blocks_conv = nn.Sequential(
                Resblock(channels=out_channels, hidden_channels=out_channels//2),
                BasicConv(out_channels, out_channels, 1)
            )
            self.concat_conv = BasicConv(out_channels*2, out_channels, 1)
        else:
            self.split_conv0 = BasicConv(out_channels, out_channels//2, 1)
            self.split_conv1 = BasicConv(out_channels, out_channels//2, 1)
            self.blocks_conv = nn.Sequential(
                *[Resblock(out_channels//2) for _ in range(num_blocks)],
                BasicConv(out_channels//2, out_channels//2, 1)
            )
            self.concat_conv = BasicConv(out_channels, out_channels, 1)

    def forward(self, x):
        x = self.downsample_conv(x)

        x0 = self.split_conv0(x)

        x1 = self.split_conv1(x)
        x1 = self.blocks_conv(x1)

        x = torch.cat([x1, x0], dim=1)
        x = self.concat_conv(x)

        return x


class CSPDarkNet(nn.Module):
    """
    相比于 YOLO_v3，主干特征提取网络 Backbone 的改进点有两个：
    1. 主干特征提取网络: DarkNet53 -> CSPDarkNet53
    2. 激活函数: 使用 Mish 激活函数
    """
    def __init__(self, layers):
        super(CSPDarkNet, self).__init__()
        self.inplanes = 32
        # DarknetConv2D_BN_Mish -> (608, 608, 32)
        self.conv1 = BasicConv(3, self.inplanes, kernel_size=3, stride=1)
        self.feature_channels = [64, 128, 256, 512, 1024]

        self.stages = nn.ModuleList([
            Resblock_body(self.inplanes, self.feature_channels[0], layers[0], first=True),  # (304, 304, 64) * 1
            Resblock_body(self.feature_channels[0], self.feature_channels[1], layers[1], first=False),  # (152, 152, 128) * 2
            Resblock_body(self.feature_channels[1], self.feature_channels[2], layers[2], first=False),  # (76, 76, 256) * 8
            Resblock_body(self.feature_channels[2], self.feature_channels[3], layers[3], first=False),  # (38, 38, 512) * 8
            Resblock_body(self.feature_channels[3], self.feature_channels[4], layers[4], first=False)   # (19, 19, 1024) * 4
        ])

        self.num_features = 1
        # 进行权值初始化
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

    def forward(self, x):
        x = self.conv1(x)

        x = self.stages[0](x)
        x = self.stages[1](x)

        # out 用于后面的特征融合（具体见网络结构图）
        out3 = self.stages[2](x)
        out4 = self.stages[3](out3)
        out5 = self.stages[4](out4)

        return out3, out4, out5


def darknet53(pretrained, **kwargs):
    model = CSPDarkNet([1, 2, 8, 8, 4])
    if pretrained:
        if isinstance(pretrained, str):
            model.load_state_dict(torch.load(pretrained))
        else:
            raise Exception("darknet request a pretrained path. got [{}]".format(pretrained))
    return model
